package ingress

import (
	"context"
	"testing"

	"time"

	awssdk "github.com/aws/aws-sdk-go-v2/aws"
	elbv2sdk "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
	elbv2types "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2/types"
	"github.com/golang/mock/gomock"
	"github.com/pkg/errors"
	"github.com/stretchr/testify/assert"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
	"sigs.k8s.io/aws-load-balancer-controller/pkg/aws/services"
	"sigs.k8s.io/aws-load-balancer-controller/pkg/model/core"
	elbv2model "sigs.k8s.io/aws-load-balancer-controller/pkg/model/elbv2"
	"sigs.k8s.io/aws-load-balancer-controller/pkg/shared_utils"
	testclient "sigs.k8s.io/controller-runtime/pkg/client/fake"
)

func Test_defaultModelBuildTask_buildAuthenticateOIDCAction(t *testing.T) {
	type env struct {
		secrets []*corev1.Secret
	}
	type args struct {
		authCfg   AuthConfig
		namespace string
	}
	authBehaviorAuthenticate := elbv2model.AuthenticateOIDCActionConditionalBehaviorAuthenticate
	tests := []struct {
		name    string
		env     env
		args    args
		want    elbv2model.Action
		wantErr error
	}{
		{
			name: "clientID & clientSecret configured",
			env: env{
				secrets: []*corev1.Secret{
					{
						ObjectMeta: metav1.ObjectMeta{
							Namespace: "my-ns",
							Name:      "my-k8s-secret",
						},
						Data: map[string][]byte{
							"clientID":     []byte("my-client-id"),
							"clientSecret": []byte("my-client-secret"),
						},
					},
				},
			},
			args: args{
				authCfg: AuthConfig{
					Type: AuthTypeCognito,
					IDPConfigOIDC: &AuthIDPConfigOIDC{
						Issuer:                "https://example.com",
						AuthorizationEndpoint: "https://authorization.example.com",
						TokenEndpoint:         "https://token.example.com",
						UserInfoEndpoint:      "https://userinfo.example.co",
						SecretName:            "my-k8s-secret",
						AuthenticationRequestExtraParams: map[string]string{
							"key1": "value1",
						},
					},
					OnUnauthenticatedRequest: "authenticate",
					Scope:                    "email",
					SessionCookieName:        "my-session-cookie",
					SessionTimeout:           65536,
				},
				namespace: "my-ns",
			},
			want: elbv2model.Action{
				Type: elbv2model.ActionTypeAuthenticateOIDC,
				AuthenticateOIDCConfig: &elbv2model.AuthenticateOIDCActionConfig{
					Issuer:                "https://example.com",
					AuthorizationEndpoint: "https://authorization.example.com",
					TokenEndpoint:         "https://token.example.com",
					UserInfoEndpoint:      "https://userinfo.example.co",
					ClientID:              "my-client-id",
					ClientSecret:          "my-client-secret",
					AuthenticationRequestExtraParams: map[string]string{
						"key1": "value1",
					},
					OnUnauthenticatedRequest: authBehaviorAuthenticate,
					Scope:                    awssdk.String("email"),
					SessionCookieName:        awssdk.String("my-session-cookie"),
					SessionTimeout:           awssdk.Int64(65536),
				},
			},
		},
		{
			name: "clientSecret has control characters at end",
			env: env{
				secrets: []*corev1.Secret{
					{
						ObjectMeta: metav1.ObjectMeta{
							Namespace: "my-ns",
							Name:      "my-k8s-secret",
						},
						Data: map[string][]byte{
							"clientID":     []byte("my-client-id"),
							"clientSecret": []byte("my-client-secret\n"),
						},
					},
				},
			},
			args: args{
				authCfg: AuthConfig{
					Type: AuthTypeCognito,
					IDPConfigOIDC: &AuthIDPConfigOIDC{
						Issuer:                "https://example.com",
						AuthorizationEndpoint: "https://authorization.example.com",
						TokenEndpoint:         "https://token.example.com",
						UserInfoEndpoint:      "https://userinfo.example.co",
						SecretName:            "my-k8s-secret",
						AuthenticationRequestExtraParams: map[string]string{
							"key1": "value1",
						},
					},
					OnUnauthenticatedRequest: "authenticate",
					Scope:                    "email",
					SessionCookieName:        "my-session-cookie",
					SessionTimeout:           65536,
				},
				namespace: "my-ns",
			},
			want: elbv2model.Action{
				Type: elbv2model.ActionTypeAuthenticateOIDC,
				AuthenticateOIDCConfig: &elbv2model.AuthenticateOIDCActionConfig{
					Issuer:                "https://example.com",
					AuthorizationEndpoint: "https://authorization.example.com",
					TokenEndpoint:         "https://token.example.com",
					UserInfoEndpoint:      "https://userinfo.example.co",
					ClientID:              "my-client-id",
					ClientSecret:          "my-client-secret",
					AuthenticationRequestExtraParams: map[string]string{
						"key1": "value1",
					},
					OnUnauthenticatedRequest: authBehaviorAuthenticate,
					Scope:                    awssdk.String("email"),
					SessionCookieName:        awssdk.String("my-session-cookie"),
					SessionTimeout:           awssdk.Int64(65536),
				},
			},
		},
		{
			name: "clientID & clientSecret configured - legacy clientId",
			env: env{
				secrets: []*corev1.Secret{
					{
						ObjectMeta: metav1.ObjectMeta{
							Namespace: "my-ns",
							Name:      "my-k8s-secret",
						},
						Data: map[string][]byte{
							"clientId":     []byte("my-client-id"),
							"clientSecret": []byte("my-client-secret"),
						},
					},
				},
			},
			args: args{
				authCfg: AuthConfig{
					Type: AuthTypeCognito,
					IDPConfigOIDC: &AuthIDPConfigOIDC{
						Issuer:                "https://example.com",
						AuthorizationEndpoint: "https://authorization.example.com",
						TokenEndpoint:         "https://token.example.com",
						UserInfoEndpoint:      "https://userinfo.example.co",
						SecretName:            "my-k8s-secret",
						AuthenticationRequestExtraParams: map[string]string{
							"key1": "value1",
						},
					},
					OnUnauthenticatedRequest: "authenticate",
					Scope:                    "email",
					SessionCookieName:        "my-session-cookie",
					SessionTimeout:           65536,
				},
				namespace: "my-ns",
			},
			want: elbv2model.Action{
				Type: elbv2model.ActionTypeAuthenticateOIDC,
				AuthenticateOIDCConfig: &elbv2model.AuthenticateOIDCActionConfig{
					Issuer:                "https://example.com",
					AuthorizationEndpoint: "https://authorization.example.com",
					TokenEndpoint:         "https://token.example.com",
					UserInfoEndpoint:      "https://userinfo.example.co",
					ClientID:              "my-client-id",
					ClientSecret:          "my-client-secret",
					AuthenticationRequestExtraParams: map[string]string{
						"key1": "value1",
					},
					OnUnauthenticatedRequest: authBehaviorAuthenticate,
					Scope:                    awssdk.String("email"),
					SessionCookieName:        awssdk.String("my-session-cookie"),
					SessionTimeout:           awssdk.Int64(65536),
				},
			},
		},
		{
			name: "missing IDPConfigOIDC",
			args: args{
				authCfg: AuthConfig{
					Type:          AuthTypeCognito,
					IDPConfigOIDC: nil,
				},
			},
			wantErr: errors.New("missing IDPConfigOIDC"),
		},
		{
			name: "missing clientID",
			env: env{
				secrets: []*corev1.Secret{
					{
						ObjectMeta: metav1.ObjectMeta{
							Namespace: "my-ns",
							Name:      "my-k8s-secret",
						},
						Data: map[string][]byte{
							"clientSecret": []byte("my-client-secret"),
						},
					},
				},
			},
			args: args{
				authCfg: AuthConfig{
					Type: AuthTypeCognito,
					IDPConfigOIDC: &AuthIDPConfigOIDC{
						Issuer:                "https://example.com",
						AuthorizationEndpoint: "https://authorization.example.com",
						TokenEndpoint:         "https://token.example.com",
						UserInfoEndpoint:      "https://userinfo.example.co",
						SecretName:            "my-k8s-secret",
						AuthenticationRequestExtraParams: map[string]string{
							"key1": "value1",
						},
					},
					OnUnauthenticatedRequest: "authenticate",
					Scope:                    "email",
					SessionCookieName:        "my-session-cookie",
					SessionTimeout:           65536,
				},
				namespace: "my-ns",
			},
			wantErr: errors.New("missing clientID, secret: my-ns/my-k8s-secret"),
		},
		{
			name: "missing clientSecret",
			env: env{
				secrets: []*corev1.Secret{
					{
						ObjectMeta: metav1.ObjectMeta{
							Namespace: "my-ns",
							Name:      "my-k8s-secret",
						},
						Data: map[string][]byte{
							"clientID": []byte("my-client-id"),
						},
					},
				},
			},
			args: args{
				authCfg: AuthConfig{
					Type: AuthTypeCognito,
					IDPConfigOIDC: &AuthIDPConfigOIDC{
						Issuer:                "https://example.com",
						AuthorizationEndpoint: "https://authorization.example.com",
						TokenEndpoint:         "https://token.example.com",
						UserInfoEndpoint:      "https://userinfo.example.co",
						SecretName:            "my-k8s-secret",
						AuthenticationRequestExtraParams: map[string]string{
							"key1": "value1",
						},
					},
					OnUnauthenticatedRequest: "authenticate",
					Scope:                    "email",
					SessionCookieName:        "my-session-cookie",
					SessionTimeout:           65536,
				},
				namespace: "my-ns",
			},
			wantErr: errors.New("missing clientSecret, secret: my-ns/my-k8s-secret"),
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			k8sSchema := runtime.NewScheme()
			clientgoscheme.AddToScheme(k8sSchema)
			k8sClient := testclient.NewClientBuilder().WithScheme(k8sSchema).Build()
			for _, secret := range tt.env.secrets {
				err := k8sClient.Create(context.Background(), secret.DeepCopy())
				assert.NoError(t, err)
			}

			task := &defaultModelBuildTask{
				k8sClient: k8sClient,
			}
			got, err := task.buildAuthenticateOIDCAction(context.Background(), tt.args.namespace, tt.args.authCfg)
			if tt.wantErr != nil {
				assert.EqualError(t, err, tt.wantErr.Error())
			} else {
				assert.NoError(t, err)
				assert.Equal(t, tt.want, got)
			}
		})
	}
}

func Test_defaultModelBuildTask_buildJwtValidationAction(t *testing.T) {
	type args struct {
		jwtValidationConfig *JwtValidationConfig
	}
	tests := []struct {
		name    string
		args    args
		want    *elbv2model.Action
		wantErr error
	}{
		{
			name: "gracefully handles nil config",
			args: args{
				jwtValidationConfig: nil,
			},
			want: nil,
		},
		{
			name: "jwt validation with no additional claims",
			args: args{
				jwtValidationConfig: &JwtValidationConfig{
					JwksEndpoint: "https://issuer.example.com/.well-known/jwks.json",
					Issuer:       "https://issuer.com",
				},
			},
			want: &elbv2model.Action{
				Type: elbv2model.ActionTypeJwtValidation,
				JwtValidationConfig: &elbv2model.JwtValidationConfig{
					JwksEndpoint: "https://issuer.example.com/.well-known/jwks.json",
					Issuer:       "https://issuer.com",
				},
			},
		},
		{
			name: "jwt validation with additional claims",
			args: args{
				jwtValidationConfig: &JwtValidationConfig{
					JwksEndpoint: "https://issuer.example.com/.well-known/jwks.json",
					Issuer:       "https://issuer.com",
					AdditionalClaims: []JwtAdditionalClaim{
						{
							Format: "string-array",
							Name:   "scope",
							Values: []string{"read:api", "write:api"},
						},
						{
							Format: "single-string",
							Name:   "iat",
							Values: []string{"12456"},
						},
						{
							Format: "space-separated-values",
							Name:   "aud",
							Values: []string{"https://example.com", "https://another-site.com"},
						},
					},
				},
			},
			want: &elbv2model.Action{
				Type: elbv2model.ActionTypeJwtValidation,
				JwtValidationConfig: &elbv2model.JwtValidationConfig{
					JwksEndpoint: "https://issuer.example.com/.well-known/jwks.json",
					Issuer:       "https://issuer.com",
					AdditionalClaims: []elbv2model.JwtAdditionalClaim{
						{
							Format: "string-array",
							Name:   "scope",
							Values: []string{"read:api", "write:api"},
						},
						{
							Format: "single-string",
							Name:   "iat",
							Values: []string{"12456"},
						},
						{
							Format: "space-separated-values",
							Name:   "aud",
							Values: []string{"https://example.com", "https://another-site.com"},
						},
					},
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t1 *testing.T) {
			task := &defaultModelBuildTask{}
			got, err := task.buildJwtValidationAction(context.Background(), tt.args.jwtValidationConfig)
			if tt.wantErr != nil {
				assert.EqualError(t, err, tt.wantErr.Error())
			} else {
				assert.NoError(t, err)
				assert.Equal(t, tt.want, got)
			}
		})
	}
}

func Test_defaultModelBuildTask_buildSSLRedirectAction(t *testing.T) {
	type args struct {
		sslRedirectConfig SSLRedirectConfig
	}
	tests := []struct {
		name string
		args args
		want elbv2model.Action
	}{
		{
			name: "SSLRedirect to 443 with 301",
			args: args{
				sslRedirectConfig: SSLRedirectConfig{
					SSLPort:    443,
					StatusCode: "HTTP_301",
				},
			},
			want: elbv2model.Action{
				Type: elbv2model.ActionTypeRedirect,
				RedirectConfig: &elbv2model.RedirectActionConfig{
					Port:       awssdk.String("443"),
					Protocol:   awssdk.String("HTTPS"),
					StatusCode: "HTTP_301",
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t1 *testing.T) {
			task := &defaultModelBuildTask{}
			got := task.buildSSLRedirectAction(context.Background(), tt.args.sslRedirectConfig)
			assert.Equal(t, tt.want, got)
		})
	}
}

func Test_defaultModelBuildTask_buildForwardActionWithTargetGroupName(t *testing.T) {
	type describeTargetGroupsAsListCall struct {
		req  *elbv2sdk.DescribeTargetGroupsInput
		resp []elbv2types.TargetGroup
		err  error
	}
	type args struct {
		ingress                         ClassifiedIngress
		forwardActionConfig             ForwardActionConfig
		describeTargetGroupsAsListCalls []describeTargetGroupsAsListCall
		cache                           map[string]string
	}
	tests := []struct {
		name      string
		args      args
		want      elbv2model.Action
		wantErr   error
		wantCache map[string]string
	}{
		{
			name: "Forward to target group identified by name",
			args: args{
				forwardActionConfig: ForwardActionConfig{
					TargetGroups: []TargetGroupTuple{{TargetGroupName: awssdk.String("tg-name1")}},
				},

				describeTargetGroupsAsListCalls: []describeTargetGroupsAsListCall{
					{
						req: &elbv2sdk.DescribeTargetGroupsInput{
							Names: []string{"tg-name1"},
						},
						resp: []elbv2types.TargetGroup{
							{
								TargetGroupArn: awssdk.String("tg-arn1"),
								TargetType:     elbv2types.TargetTypeEnum("instance"),
							},
						},
					},
				},
				cache: map[string]string{},
			},
			want: elbv2model.Action{
				Type: elbv2model.ActionTypeForward,
				ForwardConfig: &elbv2model.ForwardActionConfig{
					TargetGroups: []elbv2model.TargetGroupTuple{{
						TargetGroupARN: core.LiteralStringToken("tg-arn1"),
					}},
				},
			},
			wantCache: map[string]string{
				"tg-name1": "tg-arn1",
			},
		},
		{
			name: "Forward to target group identified by name is cached",
			args: args{
				forwardActionConfig: ForwardActionConfig{
					TargetGroups: []TargetGroupTuple{{TargetGroupName: awssdk.String("tg-name2")}},
				},

				describeTargetGroupsAsListCalls: []describeTargetGroupsAsListCall{},
				cache: map[string]string{
					"tg-name2": "tg-arn2",
				},
			},
			want: elbv2model.Action{
				Type: elbv2model.ActionTypeForward,
				ForwardConfig: &elbv2model.ForwardActionConfig{
					TargetGroups: []elbv2model.TargetGroupTuple{{
						TargetGroupARN: core.LiteralStringToken("tg-arn2"),
					}},
				},
			},
			wantCache: map[string]string{
				"tg-name2": "tg-arn2",
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t1 *testing.T) {
			ctrl := gomock.NewController(t)
			defer ctrl.Finish()

			elbv2Client := services.NewMockELBV2(ctrl)
			for _, call := range tt.args.describeTargetGroupsAsListCalls {
				elbv2Client.EXPECT().DescribeTargetGroupsAsList(gomock.Any(), call.req).Return(call.resp, call.err)
			}
			task := &defaultModelBuildTask{
				elbv2Client:                elbv2Client,
				targetGroupNameToArnMapper: shared_utils.NewTargetGroupNameToArnMapper(elbv2Client),
			}

			for targetGroupName, cachedArn := range tt.args.cache {
				task.targetGroupNameToArnMapper.GetCache().Set(targetGroupName, cachedArn, 10*time.Minute)
			}

			got, err := task.buildForwardAction(context.Background(), tt.args.ingress, Action{
				Type:          ActionTypeForward,
				ForwardConfig: &tt.args.forwardActionConfig,
			})
			assert.Equal(t, tt.want, got)
			assert.Equal(t, tt.wantErr, err)
			assert.Equal(t, len(tt.wantCache), task.targetGroupNameToArnMapper.GetCache().Len())
			for targetGroupName, expectedArn := range tt.wantCache {
				rawCacheItem, exists := task.targetGroupNameToArnMapper.GetCache().Get(targetGroupName)
				assert.True(t, exists)
				cachedArn := rawCacheItem.(string)
				assert.Equal(t, expectedArn, cachedArn)
			}
		})
	}
}
